/*
 * Copyright (c) GoKeep by gokeep.org. 2022-2023. All rights reserved
 */
package org.gokeep.application.config;

import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.gokeep.common.util.UserContext;
import org.gokeep.common.component.TokenComponent;
import org.gokeep.common.config.properties.AuthProperties;
import org.gokeep.common.dto.TokenPayload;
import org.gokeep.common.dto.UserContextInfo;
import org.gokeep.common.exception.GoKeepException;
import org.reactivestreams.Publisher;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.util.CollectionUtils;
import org.springframework.web.reactive.config.EnableWebFlux;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import reactor.core.publisher.Mono;

import java.util.List;

/**
 * @author xuning
 */
@Slf4j
@Configuration
@EnableWebFlux
public class WebFilterConfig implements WebFilter {

    @Resource
    private TokenComponent tokenComponent;

    @Resource
    private AuthProperties authProperties;

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
        ServerHttpRequest serverHttpRequest = exchange.getRequest();
        ServerHttpResponse serverHttpResponse = exchange.getResponse();
        String method = serverHttpRequest.getMethod().name();
        String path = serverHttpRequest.getPath().value();
        log.info("Request: {} -> {}", method, path);
        HttpHeaders headers = serverHttpRequest.getHeaders();
        if (!authProperties.isEnable()) {
            return chain.filter(exchange);
        }
        if (authProperties.getWhiteList().contains(path)) {
            return chain.filter(exchange);
        }
        List<String> authorizations = headers.get("Authorization");
        if (CollectionUtils.isEmpty(authorizations)) {
            log.warn("Request: {} -> {} 'Authorizations' is empty", method, path);
            return chain.filter(exchange.mutate().response(noAuthDecorator(serverHttpResponse)).build());
        }
        String token = authorizations.get(0);
        if (StringUtils.isNotBlank(token)) {
            try {
                TokenPayload tokenPayload = tokenComponent.verifyTokenGetPayload(token);
                UserContextInfo userContextInfo = new UserContextInfo();
                userContextInfo.setToken(token).setUserId(tokenPayload.getUserId()).setUsername(tokenPayload.getUsername()).setClientType(tokenPayload.getClientType());
                UserContext.put(userContextInfo);
                return chain.filter(exchange);
            }catch (Exception e) {
                return chain.filter(exchange.mutate().response(noAuthDecorator(serverHttpResponse)).build());
            }
        }
        log.warn("Request: {} -> {} 'Authorizations' is empty", method, path);
        return chain.filter(exchange.mutate().response(noAuthDecorator(serverHttpResponse)).build());
    }

    private ServerHttpResponseDecorator noAuthDecorator(ServerHttpResponse original) {
        return serverHttpResponseDecorator(original, HttpStatus.UNAUTHORIZED, new GoKeepException(HttpStatus.UNAUTHORIZED, "Not Authorized", "未授权"));
    }

    private ServerHttpResponseDecorator serverHttpResponseDecorator(ServerHttpResponse original, HttpStatus status, GoKeepException e) {
        ServerHttpResponseDecorator decorator = new ServerHttpResponseDecorator(original) {
            @Override
            public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
                Mono<DataBuffer> mono = DataBufferUtils.join(body)
                        .map(buffer -> {
                            try {
                                byte[] bytes = e.toByteArray();
                                getHeaders().setContentLength(bytes.length);
                                return this.bufferFactory().wrap(bytes);
                            } finally {
                                DataBufferUtils.release(buffer);
                            }
                        });
                return super.writeWith(mono);
            }
        };
        decorator.setStatusCode(status);
        return decorator;
    }
}